import sys

from PySide6.QtCore import Qt
from PySide6.QtGui import QIcon, QKeySequence
from PySide6.QtWidgets import QApplication, QWidget, QTableWidget, QTableWidgetItem, QMenu, QVBoxLayout, QMessageBox


class twContextDemo(QWidget):
    def __init__(self):
        super().__init__()
        self.initUI()

    def initUI(self):
        self.setGeometry(300, 300, 460, 360)
        self.setWindowTitle('QTableWidget扩展表格显示上下文菜单的演示')
        # self.setWindowIcon(QIcon('../web.ico'))

        layout = QVBoxLayout()
        self.tablewidget = QTableWidget()
        self.tablewidget.setRowCount(3)
        self.tablewidget.setColumnCount(3)
        layout.addWidget(self.tablewidget)
        self.tablewidget.setHorizontalHeaderLabels(['姓名', '性别', '体重(kg)'])

        self.addSampleData()

        # 允许打开上下文菜单
        self.tablewidget.setContextMenuPolicy(Qt.CustomContextMenu)
        # 绑定事件
        self.tablewidget.customContextMenuRequested.connect(self.generateMenu)

        self.setLayout(layout)

    def addSampleData(self):
        data = [
            ['美国队长', '男', '180'],
            ['拿破仑', '女', '150'],
            ['黑寡妇', '女', '75']
        ]
        for row, rowData in enumerate(data):
            for col, value in enumerate(rowData):
                item = QTableWidgetItem(value)
                self.tablewidget.setItem(row, col, item)

    def generateMenu(self, pos):
        # 获取点击行号
        selectedIndexes = self.tablewidget.selectionModel().selectedIndexes()
        if len(selectedIndexes) > 0:
            rowNum = selectedIndexes[0].row()
            if rowNum < 2:  # 如果选择的行索引小于2，弹出上下文菜单
                menu = QMenu()
                item1 = menu.addAction("复制", self.onCopy, QKeySequence.Copy)
                item2 = menu.addAction("剪切", self.onCut, QKeySequence.Cut)
                item3 = menu.addAction("粘贴", self.onPaste, QKeySequence.Paste)
                item4 = menu.addAction("添加一行", self.addRow)
                item5 = menu.addAction("删除", self.onDelete)
                item6 = menu.addAction("修改", self.onEdit)

                # 转换坐标系
                screenPos = self.tablewidget.mapToGlobal(pos)

                # 显示菜单
                action = menu.exec(screenPos)
                if action:
                    self.handleMenuAction(action, rowNum)

    def handleMenuAction(self, action, rowNum):
        text = action.text()
        if text == "复制":
            self.onCopy(rowNum)
        elif text == "剪切":
            self.onCut(rowNum)
        elif text == "粘贴":
            self.onPaste(rowNum)
        elif text == "添加一行":
            self.addRow(rowNum)
        elif text == "删除":
            self.onDelete(rowNum)
        elif text == "修改":
            self.onEdit(rowNum)

    def onCopy(self, rowNum):
        item1 = self.tablewidget.item(rowNum, 0)
        item2 = self.tablewidget.item(rowNum, 1)
        item3 = self.tablewidget.item(rowNum, 2)
        print(f'选择了复制，行号：{rowNum}, 数据：{item1.text()}, {item2.text()}, {item3.text()}')

    def onCut(self, rowNum):
        print(f'选择了剪切，行号：{rowNum}')

    def onPaste(self, rowNum):
        print(f'选择了粘贴，行号：{rowNum}')

    def addRow(self, rowNum):
        self.tablewidget.insertRow(rowNum + 1)
        print(f'在行号 {rowNum + 1} 处添加了一行')

    def onDelete(self, rowNum):
        if QMessageBox.question(self, "确认", "确定要删除选定行吗？", QMessageBox.Yes | QMessageBox.No) == QMessageBox.Yes:
            self.tablewidget.removeRow(rowNum)
            print(f'已删除行号 {rowNum}')

    def onEdit(self, rowNum):
        print(f'选择了修改，行号：{rowNum}')


if __name__ == '__main__':
    app = QApplication(sys.argv)
    # 设置应用图标
    app.setWindowIcon(QIcon('../web.ico'))
    w = twContextDemo()
    w.show()
    sys.exit(app.exec())
